{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "a28782fc",
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "# %load_ext nb_black"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "3d71b56b",
   "metadata": {},
   "outputs": [],
   "source": [
    "from bert_entity_linker import surface_index_memory,BertEntityLinker\n",
    "from common_utils import read_jsonl,read_json,save_to_json\n",
    "from tqdm import tqdm\n",
    "from copy import deepcopy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "452f49d2",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "loading: data/freebase_complete_all_mention_mid_vocab\n",
      "loading: data/freebase_complete_all_mention_surface_index\n",
      "loading: data/entity_list_file_freebase_complete_all_mention\n"
     ]
    }
   ],
   "source": [
    "surface_index = surface_index_memory.EntitySurfaceIndexMemory(\n",
    "    \"data/entity_list_file_freebase_complete_all_mention\",\n",
    "    \"data/surface_map_file_freebase_complete_all_mention\",\n",
    "    \"data/freebase_complete_all_mention\",\n",
    ")\n",
    "entity_linker = BertEntityLinker(\n",
    "    surface_index, model_path=\"/data/trained_ner_model/\", device=\"cuda:0\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "2c10c03e",
   "metadata": {},
   "outputs": [],
   "source": [
    "x = entity_linker.identify_entities(\n",
    "    \"safety and tolerance of intermittent intravenous and oral zidovudine therapy in human immunodeficiency virus-infected pediatric patients. pediatric zidovudine phase i study group. is a medical trial for what?\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "4eb97abb",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'m.040fvq9'"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "x[0].sparql_name()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "f9711ea9",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<aaqu_entity_linker.KBEntity at 0x7efe6001d5e0>"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "x[0].entity"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "df604c6e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'name': 'safety and tolerance of intermittent intravenous and oral zidovudine therapy in human immunodeficiency virus-infected pediatric patients. pediatric zidovudine phase i study group.',\n",
       " 'mention': 'safety and tolerance of intermittent intravenous and oral zidovudine therapy in human immunodeficiency virus-infected pediatric patients . pediatric zidovudine phase i study group .',\n",
       " 'surface_score': 1.0,\n",
       " 'score': 1,\n",
       " 'sparql_name': 'm.040fvq9',\n",
       " 'perfect_match': True}"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "x[0].to_dict"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "0421d38a",
   "metadata": {},
   "outputs": [],
   "source": [
    "def el(id,text):\n",
    "    res = entity_linker.identify_entities(text)\n",
    "    res = [r.to_dict for r in res]\n",
    "    return {\"id\":id, \"el\":deepcopy(res)}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "f632d589",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'id': '123',\n",
       " 'el': [{'name': 'iphone 6',\n",
       "   'mention': 'iphone 6',\n",
       "   'surface_score': 0.3333333333333333,\n",
       "   'score': 1,\n",
       "   'sparql_name': 'm.011lbbq1',\n",
       "   'perfect_match': True},\n",
       "  {'name': 'apple iphone 6',\n",
       "   'mention': 'iphone 6',\n",
       "   'surface_score': 0.3333333333333333,\n",
       "   'score': 1,\n",
       "   'sparql_name': 'm.011sm55g',\n",
       "   'perfect_match': False},\n",
       "  {'name': 'iphone 6',\n",
       "   'mention': 'iphone 6',\n",
       "   'surface_score': 0.3333333333333333,\n",
       "   'score': 1,\n",
       "   'sparql_name': 'm.0n_j7lc',\n",
       "   'perfect_match': True}]}"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "el(\"123\",\"iphone 6 is a wonderful product!\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f0b473a3",
   "metadata": {},
   "source": [
    "## grailqa"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "867eb12a",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "64331"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "grailqas = []\n",
    "ids = {}\n",
    "for p in [\"/home/jimx/codes/rng-kbqa/GrailQA/outputs/grailqa_v1.0_train.json\",\n",
    "         \"/home/jimx/codes/rng-kbqa/GrailQA/outputs/grailqa_v1.0_dev.json\",\n",
    "         \"/home/jimx/codes/rng-kbqa/GrailQA/outputs/grailqa_v1.0_test.json\"]:\n",
    "    for item in read_json(p):\n",
    "        _id = item[\"qid\"]\n",
    "        assert _id not in ids\n",
    "        ids[_id] = 1\n",
    "        tmp = {\"id\":_id,\"text\":item[\"question\"].lower()}\n",
    "        grailqas.append(tmp)\n",
    "len(grailqas)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "1652e163",
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 64331/64331 [12:37<00:00, 84.92it/s]\n"
     ]
    }
   ],
   "source": [
    "preds = [el(**i) for i in tqdm(grailqas[:])]"
   ]
  },
  {
   "cell_type": "raw",
   "id": "358f0657",
   "metadata": {},
   "source": [
    "preds"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "fe0149b6",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "SAVING: grailqa_ents_BERTNER.json\n",
      "huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...\n",
      "To disable this warning, you can either:\n",
      "\t- Avoid using `tokenizers` before the fork if possible\n",
      "\t- Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)\n",
      "-rw-rw-r-- 1 jimx jimx 22M 5月  26 16:47 grailqa_ents_BERTNER.json\n",
      "\n"
     ]
    }
   ],
   "source": [
    "save_to_json(preds,\"grailqa_ents_BERTNER.json\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d5fd57fb",
   "metadata": {},
   "source": [
    "## wqcwq"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "adfa265a",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "34689"
      ]
     },
     "execution_count": 24,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "wqcwqs = []\n",
    "ids = {}\n",
    "for p in [\n",
    "    \"/home/jimx/codes/KnowledgeGraphBasic/ComplexWebQuestions/ComplexWebQuestions_train.json\",\n",
    "    \"/home/jimx/codes/KnowledgeGraphBasic/ComplexWebQuestions/ComplexWebQuestions_dev.json\",\n",
    "    \"/home/jimx/codes/KnowledgeGraphBasic/ComplexWebQuestions/ComplexWebQuestions_test.json\",\n",
    "    ]:\n",
    "    for item in read_json(p):\n",
    "        _id = item[\"ID\"]\n",
    "        assert _id not in ids\n",
    "        ids[_id] = 1\n",
    "        tmp = {\"id\":_id,\"text\":item[\"question\"].lower()}\n",
    "        wqcwqs.append(tmp)\n",
    "len(wqcwqs)"
   ]
  },
  {
   "cell_type": "raw",
   "id": "5dd8da91",
   "metadata": {},
   "source": [
    "el(\"123\",\"what is the state of the government in the country with government officer named nguyá»",
    "n há»¯u thá»?\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "da5a60e4",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 34689/34689 [06:58<00:00, 82.99it/s]\n"
     ]
    }
   ],
   "source": [
    "wqcwq_preds = [el(**i) for i in tqdm(wqcwqs)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "c6b51438",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "34689"
      ]
     },
     "execution_count": 27,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(wqcwq_preds)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "1c427daa",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'id': 'WebQTrn-3513_7c4117891abf63781b892537979054c6',\n",
       " 'el': [{'name': 'george washington colonials men$0027s basketball',\n",
       "   'mention': \"george washington colonials men 's basketball\",\n",
       "   'surface_score': 1.0,\n",
       "   'score': 1,\n",
       "   'sparql_name': 'm.03d0l76',\n",
       "   'perfect_match': False}]}"
      ]
     },
     "execution_count": 29,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "wqcwq_preds[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "b77a74ce",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "SAVING: wqcwq_ents_BERTNER.json\n",
      "huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...\n",
      "To disable this warning, you can either:\n",
      "\t- Avoid using `tokenizers` before the fork if possible\n",
      "\t- Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)\n",
      "-rw-rw-r-- 1 jimx jimx 15M 5月  26 17:04 wqcwq_ents_BERTNER.json\n",
      "\n"
     ]
    }
   ],
   "source": [
    "save_to_json(wqcwq_preds,\"wqcwq_ents_BERTNER.json\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4b6e0017",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "d749c1aa",
   "metadata": {},
   "source": [
    "# 处理ELQ的\n",
    "\n",
    "ELQ返回的是wikiid和实体名称，没有fbid。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "id": "4da46ef1",
   "metadata": {},
   "outputs": [],
   "source": [
    "elq_grailqa = read_json(\"/home/jimx/codes/BLINK/elq/grailqa_ents.json\")\n",
    "elq_wqcwq = read_json(\"/home/jimx/codes/BLINK/elq/wqcwq_ents.json\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 52,
   "id": "ee40bee3",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'id': 2101535001000,\n",
       " 'text': 'oxybutynin chloride 5 extended release film coated tablet is the ingredients of what routed drug?',\n",
       " 'scores': [-2.906494140625, -4.025367736816406],\n",
       " 'pred_tuples_string': [['Oxybutynin', 'oxybutynin chloride'],\n",
       "  ['This Is Your Brain on Drugs', 'what routed drug']],\n",
       " 'pred_triples': [['876274', 0, 6], ['606774', 16, 19]],\n",
       " 'tokens': [23060,\n",
       "  2100,\n",
       "  8569,\n",
       "  25680,\n",
       "  2378,\n",
       "  19057,\n",
       "  1019,\n",
       "  3668,\n",
       "  2713,\n",
       "  2143,\n",
       "  15026,\n",
       "  13855,\n",
       "  2003,\n",
       "  1996,\n",
       "  12760,\n",
       "  1997,\n",
       "  2054,\n",
       "  19578,\n",
       "  4319,\n",
       "  1029],\n",
       " 'wikiids': ['Q1060922', 'Q7785995']}"
      ]
     },
     "execution_count": 52,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "elq_grailqa[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 77,
   "id": "fc323d22",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 64331/64331 [00:37<00:00, 1732.34it/s]\n"
     ]
    }
   ],
   "source": [
    "for item in tqdm(elq_grailqa):\n",
    "    nermatchs = []\n",
    "    for j in item[\"pred_tuples_string\"]:\n",
    "        res = entity_linker.surface_index.get_entities_for_surface(j[0].lower())\n",
    "        if res:\n",
    "            nermatchs.append(res[0][0].id)\n",
    "        else:\n",
    "            nermatchs.append(None)\n",
    "    item[\"grailNER\"] = nermatchs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 78,
   "id": "783d9026",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "SAVING: /home/jimx/codes/BLINK/elq/grailqa_ents.json\n",
      "huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...\n",
      "To disable this warning, you can either:\n",
      "\t- Avoid using `tokenizers` before the fork if possible\n",
      "\t- Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)\n",
      "-rw-rw-r-- 1 jimx jimx 63M 5月  26 14:09 /home/jimx/codes/BLINK/elq/grailqa_ents.json\n",
      "\n"
     ]
    }
   ],
   "source": [
    "save_to_json(elq_grailqa,(\"/home/jimx/codes/BLINK/elq/grailqa_ents.json\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "85346871",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 79,
   "id": "94d7489a",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'id': 'WebQTrn-3513_7c4117891abf63781b892537979054c6',\n",
       " 'text': \"what state is home to the university that is represented in sports by george washington colonials men's basketball?\",\n",
       " 'scores': [-0.12578517198562622],\n",
       " 'pred_tuples_string': [[\"George Washington Colonials men's basketball\",\n",
       "   \"george washington colonials men's basketball\"]],\n",
       " 'pred_triples': [['1925359', 13, 21]],\n",
       " 'tokens': [2054,\n",
       "  2110,\n",
       "  2003,\n",
       "  2188,\n",
       "  2000,\n",
       "  1996,\n",
       "  2118,\n",
       "  2008,\n",
       "  2003,\n",
       "  3421,\n",
       "  1999,\n",
       "  2998,\n",
       "  2011,\n",
       "  2577,\n",
       "  2899,\n",
       "  5336,\n",
       "  2015,\n",
       "  2273,\n",
       "  1005,\n",
       "  1055,\n",
       "  3455,\n",
       "  1029],\n",
       " 'wikiids': ['Q2984247']}"
      ]
     },
     "execution_count": 79,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "elq_wqcwq[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 80,
   "id": "7bee1bd9",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 34689/34689 [00:27<00:00, 1281.40it/s]\n"
     ]
    }
   ],
   "source": [
    "for item in tqdm(elq_wqcwq):\n",
    "    nermatchs = []\n",
    "    for j in item[\"pred_tuples_string\"]:\n",
    "        res = entity_linker.surface_index.get_entities_for_surface(j[0].lower())\n",
    "        if res:\n",
    "            nermatchs.append(res[0][0].id)\n",
    "        else:\n",
    "            nermatchs.append(None)\n",
    "    item[\"grailNER\"] = nermatchs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 81,
   "id": "da9cf92b",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "SAVING: /home/jimx/codes/BLINK/elq/wqcwq_ents.json\n",
      "huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...\n",
      "To disable this warning, you can either:\n",
      "\t- Avoid using `tokenizers` before the fork if possible\n",
      "\t- Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)\n",
      "-rw-rw-r-- 1 jimx jimx 38M 5月  26 14:10 /home/jimx/codes/BLINK/elq/wqcwq_ents.json\n",
      "\n"
     ]
    }
   ],
   "source": [
    "save_to_json(elq_wqcwq,(\"/home/jimx/codes/BLINK/elq/wqcwq_ents.json\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 82,
   "id": "7096ebd7",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'id': 'WebQTrn-2518_1ef15e22372df70baf01b72850deb14d',\n",
       " 'text': 'the artist from the history world tour concert tour played what character on the wiz?',\n",
       " 'scores': [-2.09945011138916, -3.1760849952697754],\n",
       " 'pred_tuples_string': [['HIStory World Tour', 'history world tour'],\n",
       "  ['The Wiz', 'wiz']],\n",
       " 'pred_triples': [['827837', 4, 7], ['177102', 14, 16]],\n",
       " 'tokens': [1996,\n",
       "  3063,\n",
       "  2013,\n",
       "  1996,\n",
       "  2381,\n",
       "  2088,\n",
       "  2778,\n",
       "  4164,\n",
       "  2778,\n",
       "  2209,\n",
       "  2054,\n",
       "  2839,\n",
       "  2006,\n",
       "  1996,\n",
       "  15536,\n",
       "  2480,\n",
       "  1029],\n",
       " 'wikiids': ['Q44572', 'Q2855727'],\n",
       " 'grailNER': ['m.09x1nf', 'm.0264l0']}"
      ]
     },
     "execution_count": 82,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "elq_wqcwq[10]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a0dd1594",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.5"
  },
  "toc": {
   "base_numbering": 1,
   "nav_menu": {},
   "number_sections": true,
   "sideBar": true,
   "skip_h1_title": false,
   "title_cell": "Table of Contents",
   "title_sidebar": "Contents",
   "toc_cell": false,
   "toc_position": {},
   "toc_section_display": true,
   "toc_window_display": true
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
